/*
 * @Author: huangyupei huangyupei2021@ia.ac.cn
 * @Date: 2024-11-11 16:26:19
 * @LastEditors: huangyupei huangyupei2021@ia.ac.cn
 * @LastEditTime: 2025-03-18 22:16:36
 * @FilePath: /SelfCalib_OptiAcoustic/src/extrinsicRefinement.cc
 * @Description: 这是默认设置,请设置`customMade`, 打开koroFileHeader查看配置 进行设置: https://github.com/OBKoro1/koro1FileHeader/wiki/%E9%85%8D%E7%BD%AE
 */

// #include "acoustic_BA_g2o.h"
#include "extrinsicRefinement.h"

using namespace ExtrinRefine;


void ExtrinRefine::ExtrinsicRefineUseCam(Sophus::SE3d &T_sc, vector<Vector3d> &vPointsCamera, vector<Sophus::SE3d> vSonarPoses, vector<vector<Vector2d>> vCameraMeasAll) {
    // ~ 构建g2o优化器
    typedef g2o::BlockSolver<g2o::BlockSolverTraits<-1, -1>> BlockSolverType;
    typedef g2o::LinearSolverDense<BlockSolverType::PoseMatrixType> LinearSolverType;
    // typedef g2o::LinearSolverEigen<BlockSolverType::PoseMatrixType> LinearSolverType;

    // ~ solver 这一块的定义是参考ORB-SLAM3的
    BlockSolverType::LinearSolverType* linearSolver = new LinearSolverType();
    BlockSolverType* solver_ptr = new BlockSolverType(linearSolver);
    auto solver = new g2o::OptimizationAlgorithmLevenberg(solver_ptr);
    // auto solver = new g2o::OptimizationAlgorithmGaussNewton(solver_ptr);
    // auto solver = new g2o::OptimizationAlgorithmDogleg(solver_ptr);

    // 定义优化器
    g2o::SparseOptimizer optimizer;
    optimizer.setAlgorithm(solver);
    optimizer.setVerbose(false); 

    // 构建 g2o 问题
    // Step 1：添加顶点
    // ~ 待优化顶点：外参变量 T_sc
    VertexPose *vTsc = new VertexPose();
    vTsc->setId(0);
    vTsc->setEstimate(T_sc);
    vTsc->setFixed(false);
    optimizer.addVertex(vTsc);

    // ~ 固定顶点：声呐的相对运动 T_si_s0
    // 添加在第2~7个相机坐标系下的投影
    int num_sonarPose = vSonarPoses.size();     
    vector<VertexPose *> vertex_sonarPoses;
    for (int i = 1; i < num_sonarPose; i++) {
        VertexPose *vT_si_s0 = new VertexPose();
        vT_si_s0->setId(i);                         // Id范围：1 ~ 6
        vT_si_s0->setEstimate(vSonarPoses[i].inverse());
        vT_si_s0->setFixed(true);
        optimizer.addVertex(vT_si_s0);
        vertex_sonarPoses.emplace_back(vT_si_s0);       // index范围：0 ~ 5
    }
    

    // ~ 待优化顶点：相机三角化的三维点    
    int num_points = vPointsCamera.size();
    vector<VertexPoint *> vertex_points;
    for (int i = 0; i < num_points; i++) {
        VertexPoint *vPoint = new VertexPoint();
        vPoint->setId(i + num_sonarPose);           // Id从7开始
        vPoint->setEstimate(vPointsCamera[i]);
        vPoint->setMarginalized(true);
        // vPoint->setFixed(true);             // 尝试将其固定
        optimizer.addVertex(vPoint);
        vertex_points.emplace_back(vPoint);
    }

    // Step 2: 添加重投影误差边
    // 自由度为2（单相机/声呐，测量为2维）的卡方分布，显著性水平为0.05，对应的临界阈值5.991
    // 可以理解为卡方值高于5.991 95%的几率为外点
    const float deltaMono = sqrt(5.991);
    for (int i = 2; i <= 7; i++) {      // 从第2幅相机图像到第7幅相机图像，总共有6次观测，每次观测有20个误差边
        // 读取相机图像特征点测量
        vector<Eigen::Vector2d> vCameraMeasi = vCameraMeasAll[i - 1];
        
        // 获取到相机图像特征点测量后，构建重投影误差边
        for (int j = 0; j < vCameraMeasi.size(); j++) {
            EdgeCameraReproj *e = new EdgeCameraReproj();
            e->setVertex(0, vTsc);
            e->setVertex(1, vertex_sonarPoses[i - 2]);     // 第2幅相机图像，对应顶点是vertex_sonarPoses[0]
            e->setVertex(2, vertex_points[j]);             // 该误差边观测到的3D点
            e->setMeasurement(vCameraMeasi[j]);
            // const float invSigma2 = 1 / (1e4);
            e->setInformation(Eigen::Matrix2d::Identity());

            auto rk = new g2o::RobustKernelHuber();
            e->setRobustKernel(rk);
            rk->setDelta(deltaMono);

            optimizer.addEdge(e);
        }
    }

    // Step 3: 执行优化
    optimizer.initializeOptimization();
    optimizer.optimize(40);

    // Step 4: 取出优化的结果
    T_sc = vTsc->estimate();
    
    for (int i = 0; i < num_points; i++) {
        vPointsCamera[i] = vertex_points[i]->estimate();
    }
}


void ExtrinRefine::ExtrinsicRefineUseSonar(
    Sophus::SE3d &T_sc, vector<Vector3d> &vPointsSonar, 
    vector<Sophus::SE3d> vCameraPoses, vector<vector<Vector2d>> vSonarMeasAll) {
    // ~ 构建g2o优化器
    typedef g2o::BlockSolver<g2o::BlockSolverTraits<-1, -1>> BlockSolverType;
    typedef g2o::LinearSolverDense<BlockSolverType::PoseMatrixType> LinearSolverType;
    // typedef g2o::LinearSolverEigen<BlockSolverType::PoseMatrixType> LinearSolverType;

    // ~ solver 这一块的定义参考ORB-SLAM3
    BlockSolverType::LinearSolverType* linearSolver = new LinearSolverType();
    BlockSolverType* solver_ptr = new BlockSolverType(linearSolver);
    auto solver = new g2o::OptimizationAlgorithmLevenberg(solver_ptr);

    // 定义优化器
    g2o::SparseOptimizer optimizer;
    optimizer.setAlgorithm(solver);
    optimizer.setVerbose(false); 

    // 构建 g2o 问题
    // Step 1：添加顶点
    // ~ 待优化顶点：外参变量T_sc
    VertexPose *vTsc = new VertexPose();
    vTsc->setId(0);
    vTsc->setEstimate(T_sc);
    vTsc->setFixed(false);
    optimizer.addVertex(vTsc);

    // ~ 待优化顶点：相机的位姿(表示在第一帧相机坐标系下)
    int num_camPose = vCameraPoses.size();   
    // num_camPose = 3;              
    vector<VertexPose *> vertex_camPoses;
    for (int i = 1; i < num_camPose; i++) {
        VertexPose *vT_ci_c0 = new VertexPose();
        vT_ci_c0->setId(i);                         // Id范围：1 ~ 6
        vT_ci_c0->setEstimate(vCameraPoses[i].inverse());
        vT_ci_c0->setFixed(false);      // 相机位姿是否固定
        optimizer.addVertex(vT_ci_c0);
        vertex_camPoses.emplace_back(vT_ci_c0);       // index范围：0 ~ 5
    }

    // ~ 固定顶点：声呐BA恢复的三维点(表示在第一帧声呐坐标系下)
    int num_points = vPointsSonar.size();
    vector<VertexPoint *> vertex_points;
    for (int i = 0; i < num_points; i++) {
        VertexPoint *vPoint = new VertexPoint();
        vPoint->setId(i + num_camPose);           // Id从7 / num_camPose开始
        vPoint->setEstimate(vPointsSonar[i]);
        vPoint->setFixed(true);
        optimizer.addVertex(vPoint);
        vertex_points.emplace_back(vPoint);
    }
    

    // Step 2: 添加重投影误差边
    // 自由度为2的卡方分布，显著性水平为0.05，对应的临界阈值5.991
    // 可以理解为卡方值高于5.991 95%的几率为外点
    const float deltaMono = sqrt(5.991);
    
    for (int i = 1; i < num_camPose; i++) {      // 从第2幅图像到第7幅图像，总共有6次观测，每次观测有20个误差边
        // 读取声呐图像特征点测量
        vector<Eigen::Vector2d> vSonarMeasi = vSonarMeasAll[i];
        
        // 获取到声呐图像特征点测量后，构建重投影误差边
        for (int j = 0; j < vSonarMeasi.size(); j++) {
            EdgeSonarReproj *e = new EdgeSonarReproj();
            e->setVertex(0, vTsc);
            e->setVertex(1, vertex_camPoses[i - 1]);            // 第2幅声呐图像，对应顶点vertex_camPoses[0]
            e->setVertex(2, vertex_points[j]);                  // 该误差边观测到的3D点
            e->setMeasurement(vSonarMeasi[j]);
            e->setInformation(Eigen::Matrix2d::Identity());

            auto rk = new g2o::RobustKernelHuber();
            e->setRobustKernel(rk);
            rk->setDelta(deltaMono);

            optimizer.addEdge(e);
        }
    }

    // Step 3: 执行优化
    optimizer.initializeOptimization();
    optimizer.optimize(40);

    // Step 4: 取出优化的结果
    T_sc = vTsc->estimate();
    
    for (int i = 1; i < num_camPose; i++) {
        vCameraPoses[i] = vertex_camPoses[i-1]->estimate().inverse();
        // cout << "vCameraPoses[" << i << "]:\n" << vCameraPoses[i].matrix() << endl;
    }
}


void ExtrinRefine::TrajConsistencyOptimization(const vector<Sophus::SE3d>& vCameraPoses, const vector<Sophus::SE3d>& vSonarPoses, Sophus::SE3d& T_sc) {
    // //~ 创建优化器[写法来自GPT]
    // g2o::SparseOptimizer optimizer;
    // auto linearSolver = g2o::make_unique<g2o::LinearSolverDense<g2o::BlockSolverX::PoseMatrixType>>();
    // auto solver = g2o::make_unique<g2o::OptimizationAlgorithmLevenberg>(
    //     g2o::make_unique<g2o::BlockSolverX>(std::move(linearSolver))
    // );
    // optimizer.setAlgorithm(solver.get());

    //~ 创建优化器
    typedef g2o::BlockSolver<g2o::BlockSolverTraits<6, 3>> BlockSolverType;     // pose is 6, landmark is 3
    typedef g2o::LinearSolverDense<BlockSolverType::PoseMatrixType> LinearSolverType;    // 线性求解器类型
    
    // 梯度下降方法,可以从GN LM DogLeg选
    BlockSolverType::LinearSolverType* linearSolver = new LinearSolverType();
    BlockSolverType* solver_ptr = new BlockSolverType(linearSolver);
    auto solver = new g2o::OptimizationAlgorithmLevenberg(solver_ptr);

    g2o::SparseOptimizer optimizer;
    optimizer.setAlgorithm(solver);
    optimizer.setVerbose(false); 

    //~ 添加优化变量（T_sc）
    g2o::VertexSE3Expmap* v_Tsc = new g2o::VertexSE3Expmap();
    // v_Tsc->setEstimate(g2o::SE3Quat(T_sc.rotationMatrix(), T_sc.translation()));
    Sophus::SE3d T = Sophus::SE3d();        // 单位阵
    v_Tsc->setEstimate(g2o::SE3Quat(T.unit_quaternion(), T.translation()));
    v_Tsc->setId(0);
    v_Tsc->setFixed(false);
    optimizer.addVertex(v_Tsc);

    //~ 添加误差项
    for (size_t i = 0; i < vCameraPoses.size(); i++) {
        EdgePoseAlignment* edge = new EdgePoseAlignment(vCameraPoses[i], vSonarPoses[i]);
        edge->setVertex(0, v_Tsc);
        Eigen::Matrix<double, 6, 6> info = Eigen::Matrix<double, 6, 6>::Identity();
        // info.topLeftCorner(3, 3) *= 1.0;  // 增强旋转部分
        // info.bottomRightCorner(3, 3) *= 1.0; // 增强平移部分
        edge->setInformation(info);  // 信息矩阵（单位矩阵）
        optimizer.addEdge(edge);
    }

    //~ 运行优化
    optimizer.initializeOptimization();
    optimizer.optimize(10);

    //~ 取出优化结果
    g2o::VertexSE3Expmap* T_sc_optimized = static_cast<g2o::VertexSE3Expmap *>(optimizer.vertex(0));
    auto SE3quat_recov = T_sc_optimized->estimate();
    T_sc = Sophus::SE3d(SE3quat_recov.rotation(), SE3quat_recov.translation());
}